import torch
import time
import unittest
from mesh_coverage_loss import _p2e_logproba

class TestP2ELogproba(unittest.TestCase):
    def setUp(self):
        torch.manual_seed(1)
        torch.cuda.manual_seed_all(1)
        self.n_s = 1
        self.n_e = 1
        
        self.sxy = torch.randn(self.n_s, 2, dtype=torch.float32).cuda()
        self.oxy = torch.randn(self.n_e, 2, dtype=torch.float32).cuda().detach().requires_grad_(True)
        theta = torch.rand(self.n_e) * torch.pi
        self.invcov = torch.stack([torch.cos(theta), -torch.sin(theta), torch.sin(theta), torch.cos(theta)], dim=1).reshape(self.n_e, 2, 2).contiguous().cuda()
        self.logdet_invcov = 2 * torch.logdet(self.invcov)
        

    def test_forward_shape(self):
        sxy = self.sxy
        oxy = self.oxy  
        invcov = self.invcov
        logdet_invcov = self.logdet_invcov 
        output1 = _p2e_logproba.apply(sxy, oxy, invcov, logdet_invcov)

    def test_backward_gradcheck(self):
        print(self.sxy)
        print(self.invcov)
        sxy = self.sxy
        oxy = self.oxy  
        invcov = self.invcov
        logdet_invcov = self.logdet_invcov 
        loss_a = _p2e_logproba.apply(sxy, oxy, invcov, logdet_invcov).mean()
        loss_a.backward()
        print(loss_a)
        print(self.oxy.grad)
    
    def tearDown(self):
        print("\n")


if __name__ == '__main__':
    unittest.main()